import torch
import torch.nn as nn
from torch.nn import functional as F
import math
from algorithms.utils.util import init
from algorithms.utils.transformer_act import discrete_autoregressive_act
from algorithms.utils.transformer_act import discrete_parallel_act

"""
多智能体Transformer (MAT) - 36维状态处理版本

处理流程：
1. 输入：36维状态向量 (batch, n_agent, 36)


2. 特征提取：直接映射为特征向量
   - 输入：(batch, n_agent, 36)
   - 映射：(batch, n_agent, n_embd)
   - 输出：(batch, n_agent, n_embd)  # 每个智能体一个特征向量

3. Transformer处理：智能体间注意力交互
   - 输入：(batch, n_agent, n_embd)
   - 输出：(batch, n_agent, n_embd)  # 包含智能体间交互信息

4. 价值函数：从编码后的特征计算价值
   - 输出：(batch, n_agent, 1)  # 每个智能体的价值

"""


# 初始化函数，用于权重初始化
def init_(m, gain=0.01, activate=False):
    if activate:
        # 如果激活函数为ReLU，计算合适的增益，确保为float类型
        gain = float(nn.init.calculate_gain('relu'))
        # 使用正交初始化权重，偏置初始化为0
    return init(m, nn.init.orthogonal_, lambda x: nn.init.constant_(x, 0), gain=gain)


# 自注意力机制模块
class SelfAttention(nn.Module):

    def __init__(self, n_embd, n_head, n_agent, masked=False):
        super(SelfAttention, self).__init__()

        assert n_embd % n_head == 0  # 确保嵌入维度能被头数整除
        self.masked = masked  # 是否使用掩码
        self.n_head = n_head  # 注意力头数
        # 为所有头部的key, query, value投影
        self.key = init_(nn.Linear(n_embd, n_embd))  # key线性变换
        self.query = init_(nn.Linear(n_embd, n_embd))  # query线性变换
        self.value = init_(nn.Linear(n_embd, n_embd))  # value线性变换
        # 输出投影
        self.proj = init_(nn.Linear(n_embd, n_embd))  # 输出线性变换
        # 如果使用掩码，创建因果掩码确保注意力只应用于输入序列的左侧
        self.register_buffer("mask", torch.tril(torch.ones(n_agent + 1, n_agent + 1))
                             .view(1, 1, n_agent + 1, n_agent + 1))  # 创建下三角掩码

    def forward(self, key, value, query):
        B, L, D = query.size()  # 获取batch、序列长度、嵌入维度

        # 计算所有头部的query, key, values，并将头部维度移到batch维度前面
        k = self.key(key).view(B, L, self.n_head, D // self.n_head).transpose(1, 2)  # (B, nh, L, hs)
        q = self.query(query).view(B, L, self.n_head, D // self.n_head).transpose(1, 2)  # (B, nh, L, hs)
        v = self.value(value).view(B, L, self.n_head, D // self.n_head).transpose(1, 2)  # (B, nh, L, hs)

        # 因果注意力: (B, nh, L, hs) x (B, nh, hs, L) -> (B, nh, L, L)
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))  # 计算注意力分数并缩放

        if self.masked:
            att = att.masked_fill(self.mask[:, :, :L, :L] == 0, float('-inf'))  # 应用掩码，防止信息泄露
        att = F.softmax(att, dim=-1)  # 对最后一个维度做softmax，得到注意力权重

        y = att @ v  # (B, nh, L, L) x (B, nh, L, hs) -> (B, nh, L, hs)
        y = y.transpose(1, 2).contiguous().view(B, L, D)  # 重新组装所有头部的输出

        # 输出投影
        y = self.proj(y)  # 通过线性层
        return y  # 返回注意力输出


# 编码器Transformer模块
class EncodeBlock(nn.Module):

    def __init__(self, n_embd, n_head, n_agent):
        super(EncodeBlock, self).__init__()

        self.ln1 = nn.LayerNorm(n_embd)  # 第一层归一化
        self.ln2 = nn.LayerNorm(n_embd)  # 第二层归一化

        self.attn = SelfAttention(n_embd, n_head, n_agent, masked=False)  # 编码器不使用掩码
        self.mlp = nn.Sequential(
            init_(nn.Linear(n_embd, 1 * n_embd), activate=True),  # 前馈网络第一层
            nn.GELU(),  # GELU激活函数
            init_(nn.Linear(1 * n_embd, n_embd))  # 前馈网络第二层
        )

    def forward(self, x):
        x = self.ln1(x + self.attn(x, x, x))  # 残差连接 + 归一化 + 自注意力
        x = self.ln2(x + self.mlp(x))  # 残差连接 + 归一化 + 前馈网络
        return x  # 返回输出


# 解码器Transformer块
class DecodeBlock(nn.Module):

    def __init__(self, n_embd, n_head, n_agent):
        super(DecodeBlock, self).__init__()

        self.ln1 = nn.LayerNorm(n_embd)  # 第一层归一化
        self.ln2 = nn.LayerNorm(n_embd)  # 第二层归一化
        self.ln3 = nn.LayerNorm(n_embd)  # 第三层归一化
        # 注意力机制（需要带掩码的)
        self.attn1 = SelfAttention(n_embd, n_head, n_agent, masked=True)
        self.attn2 = SelfAttention(n_embd, n_head, n_agent, masked=True)
        self.mlp = nn.Sequential(
            init_(nn.Linear(n_embd, 1 * n_embd), activate=True),  # 前馈网络第一层
            nn.GELU(),  # GELU激活函数
            init_(nn.Linear(1 * n_embd, n_embd))  # 前馈网络第二层
        )

    def forward(self, x, rep_enc):
        """
        DecodeBlock前向传播 - 实现自回归动作解码
        
        Args:
            x: 动作嵌入，形状为 (batch, n_agent, n_embd)
            rep_enc: 编码后的观测表示，形状为 (batch, n_agent, n_embd)
            
        Returns:
            x: 解码后的动作表示，形状为 (batch, n_agent, n_embd)
        """
        # 第一个掩码注意力：动作序列的自注意力
        # 输入：(batch, n_agent, n_embd)
        # 输出：(batch, n_agent, n_embd)
        x = self.ln1(x + self.attn1(x, x, x))  # 第一层残差连接 + 归一化 + 自注意力

        # 第二个掩码注意力：交叉注意力（动作查询，观测键值）
        # 这是MAT架构图中的关键部分：动作生成时利用观测信息
        # Q: 动作嵌入 (batch, n_agent, n_embd)
        # K, V: 编码后的观测 (batch, n_agent, n_embd)
        # 输出：(batch, n_agent, n_embd)
        x = self.ln2(rep_enc + self.attn2(key=x, value=x, query=rep_enc))  # 第二层残差连接 + 归一化 + 交互注意力

        # MLP处理
        # 输入：(batch, n_agent, n_embd)
        # 输出：(batch, n_agent, n_embd)
        x = self.ln3(x + self.mlp(x))  # 第三层残差连接+归一化+前馈网络

        return x  # 返回输出


# 编码器模块
class Encoder(nn.Module):

    def __init__(self, n_block, n_embd, n_head, n_agent, obs_dim=36, dropout_rate=0.1):
        super(Encoder, self).__init__()

        # 嵌入维度、智能体数量、观测维度
        self.n_embd = n_embd
        self.n_agent = n_agent
        self.obs_dim = obs_dim

        # 观测编码器：将36维状态向量映射到嵌入空间
        self.obs_encoder = nn.Sequential(
            nn.Linear(obs_dim, n_embd),  # 36维 -> n_embd
            nn.ReLU(inplace=True),
            nn.Dropout(dropout_rate),
            nn.LayerNorm(n_embd),
            nn.GELU()
        )

        # 嵌入归一化、多层编码器块、价值函数头部
        self.ln = nn.LayerNorm(n_embd)
        self.blocks = nn.Sequential(*[EncodeBlock(n_embd, n_head, n_agent) for _ in range(n_block)])

        # 价值函数MLP - 从编码后的观测计算V(S)
        # 设计：n_embd -> n_embd -> 1，符合MAT架构图
        self.value_head = nn.Sequential(
            init_(nn.Linear(n_embd, n_embd), activate=True),  # 第一层：特征映射
            nn.GELU(),  # GELU激活函数
            nn.LayerNorm(n_embd),  # 层归一化
            init_(nn.Linear(n_embd, 1))  # 第二层：输出价值函数
        )

    def forward(self, obs):
        """
        编码器前向传播 - 处理多智能体36维状态并输出价值函数和嵌入表示
        
        Args:
            obs: 形状为 (batch, n_agent, 36) 的多智能体状态张量
                其中36维状态包含：
                - 前12维：当前智能体的状态信息
                - 中间12维：敌方信息
                - 后12维：友方信息
            
        Returns:
            v_loc: 局部价值函数输出，形状为 (batch, n_agent, 1)
            rep: 智能体嵌入表示，形状为 (batch, n_agent, n_embd)
        """
        # 处理多智能体36维状态 - 每个智能体独立处理局部状态
        batch_size = obs.shape[0]

        # 通过观测编码器将36维状态映射为特征向量
        # 输入：(batch, n_agent, 36)
        # 输出：(batch, n_agent, n_embd)
        obs_embeddings = self.obs_encoder(obs)  # 将36维状态映射到嵌入空间

        # 通过多层Transformer块处理智能体间交互 - 实现多智能体协作
        # 智能体间通过注意力机制交换信息，实现协作决策
        # 输入：(batch, n_agent, n_embd)
        # 输出：(batch, n_agent, n_embd) - 包含智能体间交互信息
        rep = self.blocks(self.ln(obs_embeddings))  # 通过多层Transformer块

        # 通过价值函数头部得到局部值 - 从编码后的观测计算V(S)
        # 输入：(batch, n_agent, n_embd)
        # 输出：(batch, n_agent, 1) - 每个智能体的价值函数
        v_loc = self.value_head(rep)  # 通过价值函数头部得到局部值

        return v_loc, rep  # 返回局部值和嵌入表示


# 解码器模块
class Decoder(nn.Module):

    def __init__(self, action_dim, n_block, n_embd, n_head, n_agent):
        super(Decoder, self).__init__()  # 调用父类初始化

        self.action_dim = action_dim  # 动作维度
        self.n_embd = n_embd  # 嵌入维度

        # 离散动作编码器、多层解码器块、输出头部
        self.action_encoder = nn.Sequential(init_(nn.Linear(action_dim + 1, n_embd, bias=False), activate=True),
                                            nn.GELU())

        self.ln = nn.LayerNorm(n_embd)
        self.blocks = nn.Sequential(*[DecodeBlock(n_embd, n_head, n_agent) for _ in range(n_block)])
        self.head = nn.Sequential(init_(nn.Linear(n_embd, n_embd), activate=True), nn.GELU(), nn.LayerNorm(n_embd),
                                  init_(nn.Linear(n_embd, action_dim)))

    def forward(self, action, obs_rep, obs=None):
        """
        解码器前向传播 - 实现自回归动作生成
        
        Args:
            action: 动作输入，形状为 (batch, n_agent, action_dim)
                   - 训练时：历史动作或目标动作
                   - 推理时：自回归生成的动作序列
            obs_rep: 编码后的观测表示，形状为 (batch, n_agent, n_embd)
                    - 来自Encoder的输出，包含智能体间交互信息
                    - 作为交叉注意力的键(K)和值(V)
            obs: 原始观测（为了兼容性，但不使用）
            
        Returns:
            logit: 动作logits，形状为 (batch, n_agent, action_dim)
        """
        # 动作编码：将离散动作映射为嵌入向量
        # 输入：(batch, n_agent, action_dim) - 例如 (32, 4, 7)
        # 输出：(batch, n_agent, n_embd) - 例如 (32, 4, 256)
        action_embeddings = self.action_encoder(action)  # 对动作编码

        # 层归一化
        x = self.ln(action_embeddings)  # 形状：(batch, n_agent, n_embd)

        # 通过多层解码器块 - 实现自回归动作生成
        for block in self.blocks:
            # 每个DecodeBlock包含：
            # 1. 动作序列的自注意力（掩码）
            # 2. 动作与观测的交叉注意力（掩码）
            # 3. MLP处理
            x = block(x, obs_rep)  # 通过多层解码器块

        # 输出头部：生成动作logits
        # 输入：(batch, n_agent, n_embd)
        # 输出：(batch, n_agent, action_dim) - 例如 (32, 4, 7)
        logit = self.head(x)  # 输出头部

        return logit  # 返回动作logits


# 多智能体Transformer - 支持36维状态处理
class MultiAgentTransformer(nn.Module):

    def __init__(self, action_dim, n_agent,
                 n_block, n_embd, n_head, device=torch.device("cpu"),
                 obs_dim=36, dropout_rate=0.1):
        super(MultiAgentTransformer, self).__init__()

        # 智能体数量、动作维度、设备和数据类型
        self.n_agent = n_agent
        self.action_dim = action_dim
        self.tpdv = dict(dtype=torch.float32, device=device)
        self.device = device

        # 构建编码器 - 处理36维状态向量
        self.encoder = Encoder(n_block, n_embd, n_head, n_agent, obs_dim, dropout_rate)  # 编码器
        # 解码器使用编码器输出的嵌入维度
        self.decoder = Decoder(action_dim, n_block, n_embd, n_head, n_agent)  # 解码器
        self.to(device)  # 移动到指定设备

    def forward(self, obs, action, available_actions=None):
        """
        MAT模型前向传播 - 端到端训练
        
        Args:
            obs: 形状为 (batch, n_agent, 36) 的36维状态
            action: 形状为 (batch, n_agent, 1) 的动作
            available_actions: 形状为 (batch, n_agent, act_dim) 的可用动作掩码
            
        Returns:
            action_log: 动作对数概率
            v_loc: 局部价值函数
            entropy: 动作分布熵
        """

        # 编码器处理36维状态并输出价值函数和嵌入表示
        batch_size = obs.shape[0]  # 获取batch大小
        v_loc, obs_rep = self.encoder(obs)  # 编码器前向，得到局部值和嵌入

        # 并行采样离散动作
        action = action.long()
        action_log, entropy = discrete_parallel_act(self.decoder, obs_rep, obs, action, batch_size,
                                                    self.n_agent, self.action_dim, self.tpdv,
                                                    available_actions)
        # 返回动作logits、局部值、熵
        return action_log, v_loc, entropy

    def get_actions(self, obs, available_actions=None, deterministic=False):
        """
        获取动作 - 用于推理
        
        Args:
            obs: 形状为 (batch, n_agent, 36) 的36维状态
            available_actions: 形状为 (batch, n_agent, act_dim) 的可用动作掩码
            deterministic: 是否确定性采样
            
        Returns:
            output_action: 选择的动作
            output_action_log: 动作对数概率
            v_loc: 局部价值函数
        """

        # 编码器处理36维状态并输出价值函数和嵌入表示
        batch_size = obs.shape[0]  # 获取batch大小
        v_loc, obs_rep = self.encoder(obs)  # 编码器前向

        # 离散自回归采样
        output_action, output_action_log = discrete_autoregressive_act(self.decoder, obs_rep, obs, batch_size,
                                                                       self.n_agent, self.action_dim, self.tpdv,
                                                                       available_actions, deterministic)
        # 返回动作、动作logits、局部值
        return output_action, output_action_log, v_loc

    def get_values(self, obs):
        """
        获取价值函数 - 用于价值估计
        
        Args:
            obs: 形状为 (batch, n_agent, 36) 的36维状态
            
        Returns:
            v_tot: 总价值函数
        """

        # 编码器处理36维状态并输出价值函数
        v_tot, obs_rep = self.encoder(obs)  # 编码器前向
        return v_tot  # 返回总值
